import numpy as np

class MLDetector(object):

    def __init__(self, params):
        
        self.Nr = params['Nr']
        self.Nt = params['Nt']
        self.batch_size = params['BATCH_SIZE']
        self.C = np.expand_dims(params['CONSTELLATION_COMPLEX'],0) # [1,M]
        self.candidate = self.cartesian( [ params['CONSTELLATION'] for _ in range(2*self.Nt) ] )

    def __call__(self,x,H,y,sigma2):
        """
        最大似然检测算法
        Input:
        x:发送符号 array(shape=(batch_size, 2*Nt), dtype=float32)
        y:接收信号 array(shape=(batch_size, 2*Nr), dtype=float32)
        H:信道矩阵 array(shape=(batch_size, 2*Nr, 2*Nt), dtype=float32)
        sigma2:噪声方差 array(shape=(batch_size, 1), dtype=float32)
        Variable:
        candidate:发送符号组合 array(shape=(M^(2*Nt), 2*Nt), dtype=float32)
        Output:
        nodes:指标字典
        """
        # ML Detector
        candidate = np.transpose(self.candidate,[1,0])

        Hx = np.matmul(H, np.expand_dims(candidate, 0) ) # [b, Nr, M^(Nt)]

        distance = np.sqrt(np.sum(np.square(np.expand_dims(y,-1)-Hx),1)) # [b, M^(Nt)]

        min_index = np.argmin(distance,1) # [b,]

        xhat = candidate[:,min_index].T # [b, 2*Nt]

        xhat_complex = self.slicer(xhat)
        x_complex  = x[:,0:self.Nt] + 1j*x[:,self.Nt:self.Nt*2]

        totalnum = self.batch_size * self.Nt
        rightnum = self.equal(xhat_complex, x_complex)
        errornum = totalnum - rightnum

        nodes={
            'x_complex':x_complex,
            'xhat_complex':xhat_complex,
            'totalnum':totalnum,
            'errornum':errornum,
        }

        return nodes

    def cartesian(self, arrays, out=None):
        """
        Generate a cartesian product of input arrays.

        Parameters
        ----------
        arrays : list of array-like
            1-D arrays to form the cartesian product of.
        out : ndarray
            Array to place the cartesian product in.

        Returns
        -------
        out : ndarray
            2-D array of shape (M, len(arrays)) containing cartesian products
            formed of input arrays.

        Examples
        --------
        >>> cartesian(([1, 2, 3], [4, 5], [6, 7]))
        array([[1, 4, 6],
            [1, 4, 7],
            [1, 5, 6],
            [1, 5, 7],
            [2, 4, 6],
            [2, 4, 7],
            [2, 5, 6],
            [2, 5, 7],
            [3, 4, 6],
            [3, 4, 7],
            [3, 5, 6],
            [3, 5, 7]])

        """

        arrays = [np.asarray(x) for x in arrays]
        dtype = arrays[0].dtype

        n = np.prod([x.size for x in arrays])
        if out is None:
            out = np.zeros([n, len(arrays)], dtype=dtype)

        m = int(n / arrays[0].size)
        out[:,0] = np.repeat(arrays[0], m)
        if arrays[1:]:
            self.cartesian(arrays[1:], out=out[0:m,1:])
            for j in range(1, arrays[0].size):
                out[j*m:(j+1)*m,1:] = out[0:m,1:]

        return out

    def batch_matvec_mul(self, A, b):
        """
        矩阵A与矩阵b相乘，其中A.shape=(batch_size, Nr, Nt)
        b.shape = (batch_size, Nt)
        输出矩阵C，C.shape = (batch_size, Nr)
        """
        C = np.matmul(A, np.expand_dims(b, axis=2))

        return np.squeeze(C, -1)

    def slicer(self, x):
        """
        软符号转换为硬符号
        Input:
        x: 检测器检测后的信号 Tensor(shape=(batchsize, 2*Nt), dtype=float32)
        Output:
        xhat_complex: 符号硬估计，Tensor(shape=(batchsize, Nt)， dtype=tf.int32)
        """
        x_complex = x[:, 0: self.Nt] + 1j*x[:, self.Nt: self.Nt*2]
        x_complex = np.reshape(x_complex,[-1, 1])
        indices = np.argmin(np.abs(x_complex - self.C), axis=1).astype(np.int32)
        xhat_complex = self.C[:,indices].T
        xhat_complex = np.reshape(xhat_complex,[-1,self.Nt])
        
        return xhat_complex

    def equal(self,x,y):
        """
        Computes the numbers of elements for which x and y are equal
        """
        return np.sum(np.equal(x, y).astype(np.float32))
